{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### AutoGen Core - Distributed Agent Runtime\n",
    "\n",
    "[Distributed runtime](https://microsoft.github.io/autogen/stable//user-guide/core-user-guide/core-concepts/architecture.html#distributed-agent-runtime) is suitable for multi-process applications where agents may be implemented in different programming languages and running on different machines.\n",
    "\n",
    "A distributed runtime consists of:\n",
    "- `host servicer` - The host servicer facilitates communication between agents across workers and maintains the states of connections.\n",
    "- `multiple workers`  The workers run agents and communicate with the host servicer via gateways. They advertise to the host servicer the agents they run and manage the agents’ lifecycles."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Importing required libraries\n",
    "\n",
    "from autogen_core import AgentId, MessageContext, RoutedAgent, message_handler\n",
    "from autogen_agentchat.agents import AssistantAgent\n",
    "from autogen_agentchat.messages import TextMessage\n",
    "from autogen_ext.models.openai import OpenAIChatCompletionClient\n",
    "from autogen_ext.tools.langchain import LangChainToolAdapter\n",
    "\n",
    "from langchain_community.utilities import GoogleSerperAPIWrapper\n",
    "from langchain.agents import Tool\n",
    "from IPython.display import display, Markdown\n",
    "from dataclasses import dataclass\n",
    "\n",
    "from dotenv import load_dotenv"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Loading up environment variables\n",
    "load_dotenv(override=True)\n",
    "\n",
    "ALL_IN_ONE_WORKER = False"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Simple Message class"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "@dataclass\n",
    "class Message:\n",
    "    content: str"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Starting up the host service\n",
    "The code below starts the host service in the background and accepts worker connections on port 50051."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.runtimes.grpc import GrpcWorkerAgentRuntimeHost\n",
    "\n",
    "host = GrpcWorkerAgentRuntimeHost(address=\"localhost:50051\")\n",
    "host.start() "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Introducing internet search tool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "serper = GoogleSerperAPIWrapper()\n",
    "langchain_serper =Tool(\n",
    "                    name=\"internet_search\", \n",
    "                    func=serper.run, \n",
    "                    description=\"Useful for running internet searches\")\n",
    "autogen_serper = LangChainToolAdapter(langchain_serper)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "favouring_instruction = \"To help with a decision on whether to use AutoGen in a new AI Agent project, \\\n",
    "please research and briefly respond with reasons in favor of choosing AutoGen; the pros of AutoGen.\"\n",
    "\n",
    "opposing_instruction = \"To help with a decision on whether to use AutoGen in a new AI Agent project, \\\n",
    "please research and briefly respond with reasons against choosing AutoGen; the cons of Autogen.\"\n",
    "\n",
    "judge_instruction = \"You must make a decision on whether to use AutoGen for a project. \\\n",
    "Your research team has come up with the following reasons for and against. \\\n",
    "Based purely on the research from your team, please respond with your decision and brief rationale.\""
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Creating Agents"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class PlayerOneAgent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gemini-2.0-flash\")\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client, tools=[autogen_serper], reflect_on_tool_use=True)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        return Message(content=response.chat_message.content)\n",
    "    \n",
    "class PlayerTwoAgent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gemini-2.0-flash\")\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client, tools=[autogen_serper], reflect_on_tool_use=True)\n",
    "\n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        text_message = TextMessage(content=message.content, source=\"user\")\n",
    "        response = await self._delegate.on_messages([text_message], ctx.cancellation_token)\n",
    "        return Message(content=response.chat_message.content)\n",
    "    \n",
    "class JudgeAgent(RoutedAgent):\n",
    "    def __init__(self, name: str) -> None:\n",
    "        super().__init__(name)\n",
    "        model_client = OpenAIChatCompletionClient(model=\"gemini-2.0-flash\")\n",
    "        self._delegate = AssistantAgent(name, model_client=model_client)\n",
    "        \n",
    "    @message_handler\n",
    "    async def handle_my_message_type(self, message: Message, ctx: MessageContext) -> Message:\n",
    "        favouring_message = Message(content=favouring_instruction)\n",
    "        opposing_message = Message(content=opposing_instruction)\n",
    "        player_one = AgentId(\"player1\", \"default\")\n",
    "        player_two = AgentId(\"player2\", \"default\")\n",
    "        response1 = await self.send_message(favouring_message, player_one)\n",
    "        response2 = await self.send_message(opposing_message, player_two)\n",
    "        result = f\"## Pros of AutoGen:\\n{response1.content}\\n\\n## Cons of AutoGen:\\n{response2.content}\\n\\n\"\n",
    "        judgement = f\"{judge_instruction}\\n{result}Respond with your decision and brief explanation\"\n",
    "        message = TextMessage(content=judgement, source=\"user\")\n",
    "        response = await self._delegate.on_messages([message], ctx.cancellation_token)\n",
    "        return Message(content=result + \"\\n\\n## Decision:\\n\\n\" + response.chat_message.content)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Now we can set up the worker agent runtimes. We use `GrpcWorkerAgentRuntime`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from autogen_ext.runtimes.grpc import GrpcWorkerAgentRuntime\n",
    "\n",
    "if ALL_IN_ONE_WORKER:\n",
    "\n",
    "    worker = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker.start()\n",
    "\n",
    "    await PlayerOneAgent.register(worker, \"player1\", lambda: PlayerOneAgent(\"player1\"))\n",
    "    await PlayerTwoAgent.register(worker, \"player2\", lambda: PlayerTwoAgent(\"player2\"))\n",
    "    await JudgeAgent.register(worker, \"judge\", lambda: JudgeAgent(\"judge\"))\n",
    "\n",
    "    agent_id = AgentId(\"judge\", \"default\")\n",
    "\n",
    "else:\n",
    "\n",
    "    worker1 = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker1.start()\n",
    "    await PlayerOneAgent.register(worker1, \"player1\", lambda: PlayerOneAgent(\"player1\"))\n",
    "\n",
    "    worker2 = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker2.start()\n",
    "    await PlayerTwoAgent.register(worker2, \"player2\", lambda: PlayerTwoAgent(\"player2\"))\n",
    "\n",
    "    worker = GrpcWorkerAgentRuntime(host_address=\"localhost:50051\")\n",
    "    await worker.start()\n",
    "    await JudgeAgent.register(worker, \"judge\", lambda: JudgeAgent(\"judge\"))\n",
    "    agent_id = AgentId(\"judge\", \"default\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "response = await worker.send_message(Message(content=\"Go!\"), agent_id)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "display(Markdown(response.content))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### To stop the worker runtimes, we can call `stop()`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await worker.stop()\n",
    "if not ALL_IN_ONE_WORKER:\n",
    "    await worker1.stop()\n",
    "    await worker2.stop()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### We can call `stop()` to stop the host service."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "await host.stop()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
